# http://metagram.webreply.com/downloads/nysiis.py

# 
# NYSIIS Third generation Soundex replacement
# Rev 0.1
#
# This implementation is free for re-use, without any encumbrance
# from the implementor:
#   Sean D. True, WebReply.Com
#
# Several graceful improvements were contributed by Tim Peters
# None of the remaining bugs are Tim's fault.
#
# Reimplemented from example at: http://www.dropby.com/NYSIIS.html
# Thanks Matt Perez

import re
_DIGITre = re.compile(r"\d")
_JRSRre = re.compile(r" [JS]R$")
_VIre = re.compile(r" [VI]+$")
_VOWELWre = re.compile(r"[AEIOU]+W")
_VOWELre = re.compile(r"[AEIOU]+")
_AHHAre = re.compile(r"AH|HA")
_DUPSre = re.compile(r"([A-Z])\1+")
del re

import string
_rep = string.replace

# Translate all alpha to upper case, everything else to whitespace
_tran = [" "] * 256
for c in range(ord('A'), ord('Z')+1):
    chrc = chr(c)
    _tran[c] = chrc
    _tran[ord(string.lower(chrc))] = chrc
del c, chrc

_tran = string.join(_tran, "")

# Return a string with no leading/trailing whitespace
# and possible single embedded whitespace
def collapse(s):
    s = string.translate(s, _tran, string.punctuation)
    s = string.join(string.split(s))
    return s


_suffixMap = {
    "IX": "IC",
    "EX": "EC",
    "YE": "Y",
    "EE": "Y",
    "IE": "Y",
    "DT": "D",
    "RT": "D",
    "RD": "D",
    "NT": "N",
    "ND": "N",
}

def encode(surname, debug=0):
    if _DIGITre.search(surname):
        raise ValueError("No numerics in NYSIIS input")
    # Upper case, replace punctuation with whitespace
    surname = collapse(surname)
    # Drop troubling JR/SR suffixes
    surname = _JRSRre.sub("", surname)
    # Drop troubling III and IV
    surname = _VIre.sub("", surname)
    # Remove trailiing SZ
    while surname and surname[-1] in "SZ":
        surname = surname[:-1]
    # Initial MAC -> MC and PF -> F
    # XXX In 2.0, surname.startswith("MAC")
    if surname[:3] == "MAC":
        surname = "MC" + surname[3:]
    # XXX In 2.0, surname.startswith("PF")
    if surname[:2] == "PF":
        surname = surname[1:]
    # Change spelling on some suffixes
    # XXX This function returns None if len(surname) <= 2!
    if len(surname) > 2:
        suffix = surname[-2:]
        surname = surname[:-2] + _suffixMap.get(suffix, suffix)
        # Change 'EV' to 'EF' if not at start of name
        if len(surname) > 2 and string.find(surname[2:], "EV") > -1:
            surname = surname[:2] + _rep(surname[2:], "EV", "EF")
        # Save first vowel for later
        firstletter = surname[0]
        # Remove any 'W' that follows a vowel
        surname = _VOWELWre.sub("A", surname)
        # Replace all vowels with 'A'
        surname = _VOWELre.sub("A", surname)
        # Change GHT->GT, DG->G, PH->F
        for pattern, replacement in [("GHT","GT"),
                                     ("DG","G"),
                                     ("PH","F")]:
            surname = _rep(surname, pattern, replacement);
        # Change non-initial AH, or HA to just A
        surname = surname[0] + _AHHAre.sub("A", surname[1:])
        # Change KN->N, K->C
        surname = _rep(surname, "KN", "N")
        surname = _rep(surname, "K", "C")
        # Change non initial M->N, Q->G
        if len(surname) > 1:
            surname = surname[0] + _rep(surname[1:], "M", "N")
            surname = surname[0] + _rep(surname[1:], "Q", "G")
        for pattern, replacement in [("SH",  "S"),
                                     ("SCH", "S"),
                                     ("YW",  "Y")]:
            surname = _rep(surname, pattern, replacement)
        # If not first or last character, change 'Y' to 'A'
        if len(surname) > 2:
            first = surname[0]
            mid = surname[1:]
            last = ""
            if len(mid) > 2:
                last = mid[-1]
                mid = mid[:-1]
            surname = first + _rep(mid,"Y","A") + last
        # Change 'WR' to 'R'
        surname = _rep(surname, "WR", "R")
        # If not first character, change 'Z' to 'S'
        surname = surname[0] + _rep(surname, "Z", "S")
        if surname[-2:] == "AY":
            surname = surname[:-2] + "Y"
        # remove trailing vowels
        while surname and surname[-1] == "A":
            surname = surname[:-1]
        # remove all duplicated letters
        surname = _DUPSre.sub(r"\1", surname)
        # if first char of original surname was a vowel,
        # append it to the code
        if firstletter in "AEIOU":
            surname = firstletter + surname[1:]
        return surname

def main(a):
    return encode(a)

if __name__== "__main__":
    if collapse(" Alpha bet 123") != "ALPHA BET":
        raise SystemError("collapse Alpha bet 123")
    _test_vector = [
        ("JOHN JR", "JAN"),
        ("MACDONALD", "MCDANALD"),
        ("PFIZER", "FASAR"),
        ("ayer", "AR"),
        ("ayer jr", "AR"),
        ("ayer sr", "AR"),
        ("ayer III", "AR"),
        ("ayer IV", "AR"),
        ("wilbraham", "WALBRAN"),
        ("knuuuth", "NATH"),
        ("smithsz", "SNATH"),
        ("hoawland", "HALAN"),
        ("McA'Nulty", "MCANALTY"),
        ("fauntleroye", "FANTLARY"),
        ("Minehahaha", "MAN"),
        ("O'Banion", "OBANAN"),
        ]
    _enc = encode

    def _check_enc(name, result):
        global nerrors
        r = _enc(name)
        if r <> result:
            nerrors = nerrors + 1
            print "error in NYSIIS encoding test: name", name, \
                  "expected", result, "got", r

    nerrors = 0
    for s, r in _test_vector:
        _check_enc(s,r)
    if nerrors:
        raise SystemError("NYSIIS test failed with " +
                          `nerrors` + " errors")



